#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import logging
import uuid
import json
from typing import Dict, Any, BinaryIO, Optional, List
import tempfile
from pathlib import Path

import httpx
import aiohttp
import io

from document_processor.engines.base_engine import BaseEngine

logger = logging.getLogger(__name__)

class MinerUClient(BaseEngine):
    """MinerU 服务客户端，通过 API 调用独立的 MinerU 服务
    
    这个客户端替代了直接集成 MinerU 引擎，通过 HTTP 请求与独立的 MinerU 服务通信。
    """
    
    def __init__(self, config: Dict[str, Any]):
        """初始化 MinerU 客户端
        
        Args:
            config: 配置字典，包含 MinerU 服务的配置
        """
        super().__init__(config)
        
        # 初始化 MinerU 服务配置
        self.mineru_config = config.get("mineru_service", {})
        self.service_url = self.mineru_config.get("url", "http://localhost:8001")
        self.timeout = self.mineru_config.get("timeout", 120)  # 默认超时时间（秒）
        logger.info(f"已初始化 MinerU 客户端，服务地址: {self.service_url}")

    async def convert_to_markdown(
        self, 
        file_content: BinaryIO, 
        file_name: str,
        image_option: str = "obs",
        work_dir: Optional[str] = None,
        md_filename: Optional[str] = None
    ) -> Dict[str, Any]:
        """通过 MinerU 服务将文档转换为 Markdown 格式
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            image_option: 图片处理选项，可选值："obs"（上传到OBS）或"remove"（移除图片）
            work_dir: 工作目录路径，如果提供则使用此目录，否则创建临时目录
            md_filename: 生成的Markdown文件名（不含扩展名），用于命名图片目录
            
        Returns:
            包含以下键的字典:
                - text: 提取的纯文本
                - markdown: 转换后的Markdown文本
                - images: 图片信息列表
                - temp_dir: 临时目录路径
        """
        # 检查文件格式
        file_ext = os.path.splitext(file_name)[1].lower()
        supported_formats = ['.pdf']  # 目前仅支持PDF
        if file_ext not in supported_formats:
            raise ValueError(f"MinerU 服务目前仅支持PDF格式，不支持: {file_ext}")
            
        # 使用提供的工作目录或创建临时目录
        temp_dir = work_dir if work_dir else tempfile.mkdtemp()
        os.makedirs(temp_dir, exist_ok=True)
        
        # 生成任务ID
        task_id = str(uuid.uuid4())
        logger.info(f"创建 MinerU 处理任务: {task_id}")
        
        try:
            # 准备请求数据
            request_data = {
                "task_id": task_id,
                "options": {
                    "image_option": image_option
                }
            }
            
            # 重置文件指针位置
            file_content.seek(0)
            
            # 发送请求到 MinerU 服务
            files = {'file': (file_name, file_content)}
            form_data = {'request': json.dumps(request_data)}
            
            # 使用 httpx 进行异步请求
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                logger.info(f"发送请求到 MinerU 服务: {self.service_url}/process")
                response = await client.post(
                    f"{self.service_url}/process",
                    files=files,
                    data=form_data
                )
                
                if response.status_code != 200:
                    logger.error(f"MinerU 服务请求失败，状态码: {response.status_code}")
                    logger.error(f"响应内容: {response.text}")
                    raise Exception(f"MinerU 服务请求失败，状态码: {response.status_code}")
                
                # 解析响应结果
                result = response.json()
                logger.info(f"MinerU 服务响应成功，任务ID: {result.get('task_id')}")
                
                # 从响应中获取 Markdown 内容
                markdown_content = result.get("markdown", "")
                
                # 保存 Markdown 文件到本地
                md_name = md_filename or Path(file_name).stem
                output_file = os.path.join(temp_dir, f"{md_name}.md")
                with open(output_file, "w", encoding="utf-8") as f:
                    f.write(markdown_content)
                logger.info(f"已保存 Markdown 到本地文件: {output_file}")
                
                # 返回处理结果
                return {
                    "text": result.get("text", ""),  # 纯文本
                    "markdown": markdown_content,
                    "images": [],  # 图片已由服务端处理
                    "temp_dir": temp_dir,
                    "output_md": output_file,
                    "md_filename": md_name,
                    "process_time": result.get("process_time", 0),
                    "engine_name": result.get("engine_name", "MinerU")
                }
                
        except Exception as e:
            logger.error(f"调用 MinerU 服务失败: {str(e)}")
            import traceback
            logger.error(f"详细错误: {traceback.format_exc()}")
            raise

    async def health_check(self) -> Dict[str, Any]:
        """检查 MinerU 服务的健康状态
        
        Returns:
            包含服务状态信息的字典
        """
        try:
            async with httpx.AsyncClient(timeout=5) as client:
                response = await client.get(f"{self.service_url}/health")
                if response.status_code == 200:
                    return response.json()
                else:
                    return {"status": "error", "message": f"服务响应状态码: {response.status_code}"}
        except Exception as e:
            logger.error(f"MinerU 服务健康检查失败: {str(e)}")
            return {"status": "error", "message": f"无法连接到服务: {str(e)}"}

    async def extract_images(
        self,
        file_content: BinaryIO,
        file_name: str,
        work_dir: Optional[str] = None
    ) -> Dict[str, Any]:
        """从文档中提取图片（通过 MinerU 服务）
        
        此功能由 MinerU 服务端处理，客户端仅返回空列表
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            work_dir: 工作目录
            
        Returns:
            字典，包含提取结果
        """
        logger.warning("图片提取由 MinerU 服务端处理，客户端不实现此功能")
        return {
            "images": [],
            "temp_dir": work_dir or tempfile.mkdtemp()
        } 